package rollup

import (
	"bytes"
	"encoding/binary"
	"errors"
	"hash"
	"log"
	"math"
	"os"
	"strconv"

	"github.com/consensys/gnark-crypto/accumulator/merkletree"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	"github.com/consensys/gnark/std/accumulator/merkle"
)

var (
	hFunc = mimc.NewMiMC()
)

// Queue queue for storing the request (fixed size queue)
type Queue struct {
	listRequests chan Request
}

// NewQueue creates a new queue, batchSize is the capaciy
func NewQueue(batchSize int) Queue {
	resChan := make(chan Request, batchSize)
	var res Queue
	res.listRequests = resChan
	return res
}

type Operator struct {
	InitState  []byte
	State      []byte            // list of accounts: index ∥ nonce ∥ balance ∥ pubkeyX ∥ pubkeyY, each chunk is 256 bits
	HashState  []byte            // Hashed version of the state, each chunk is 256bits: ... ∥ H(index ∥ nonce ∥ balance ∥ pubkeyX ∥ pubkeyY)) ∥ ...
	AccountMap map[string]uint64 // hashmap of all available accounts (the key is the account.pubkey.X), the value is the index of the account in the state
	nbAccounts int               // number of accounts managed by this operator
	h          hash.Hash         // hash function used to build the Merkle Tree
	q          Queue             // queue of transfers
	// batch      int               // current number of transactions in a batch
	witnesses Circuit // witnesses for the snark cicruit
}

func NewOperator() Operator {

	res := Operator{}
	res.nbAccounts = 16
	res.h = hFunc
	res.q = NewQueue(16)
	// res.batch = 10

	res.InitState = make([]byte, SizeAccount*res.nbAccounts)
	res.State = make([]byte, SizeAccount*res.nbAccounts)
	res.HashState = make([]byte, hFunc.Size()*res.nbAccounts)
	res.AccountMap = make(map[string]uint64)

	for i := 0; i < res.nbAccounts; i++ {
		acc := NewAccount(uint64(i + 1))
		copy(res.State[i*SizeAccount:i*SizeAccount+SizeAccount], acc.Serialize())
		hFunc.Reset()
		_, _ = hFunc.Write(res.State[i*SizeAccount : i*SizeAccount+SizeAccount])
		s := hFunc.Sum([]byte{})
		copy(res.HashState[i*hFunc.Size():(i+1)*hFunc.Size()], s)
		res.AccountMap[string(acc.PubKey.X.Bytes())] = uint64(i)
	}

	return res
}

// readAccount reads the account located at index i
func (o *Operator) readAccount(index uint64) (*Account, error) {
	resAcc, err := DeserializeAccount(o.State[int(index)*SizeAccount : int(index)*SizeAccount+SizeAccount])
	if err != nil {
		return resAcc, err
	}

	return resAcc, nil
}

// updateAccount updates the state according to transfer
// numRequest is the number of the request currently handled (between 0 and batchSize)
func (o *Operator) updateState(req Request, numRequest int) error {

	var posSender uint64 // posReceiver
	var ok bool
	segmentSize := o.h.Size()

	// read sender's account
	b := req.PubKey.X.Bytes()
	if posSender, ok = o.AccountMap[string(b[:])]; !ok {
		return errors.New("")
	}
	senderAccount, err := o.readAccount(posSender)
	if err != nil {
		return err
	}

	o.witnesses.SenderAccountsBefore[numRequest].Index = senderAccount.Index
	o.witnesses.SenderAccountsBefore[numRequest].Nonce = senderAccount.Nonce

	//  Set witnesses for the proof of inclusion of sender and receivers account before update
	var buf bytes.Buffer
	_, err = buf.Write(o.HashState)
	if err != nil {
		return err
	}
	merkleRootBefore, proofInclusionSenderBefore, numLeaves, err := merkletree.BuildReaderProof(&buf, o.h, segmentSize, posSender)
	if err != nil {
		return err
	}
	merkletree.VerifyProof(o.h, merkleRootBefore, proofInclusionSenderBefore, posSender, numLeaves)
	merkleProofHelperSenderBefore := merkle.GenerateProofHelper(proofInclusionSenderBefore, posSender, numLeaves)

	o.witnesses.RootHashesBefore[numRequest] = merkleRootBefore
	for i := 0; i < len(proofInclusionSenderBefore); i++ {
		o.witnesses.MerkleProofsSenderBefore[numRequest][i] = proofInclusionSenderBefore[i]
		if i < len(merkleProofHelperSenderBefore) {
			o.witnesses.MerkleProofHelperSenderBefore[numRequest][i] = merkleProofHelperSenderBefore[i]
		}
	}

	resSig := req.Verify()
	if !resSig {
		return errors.New("")
	}

	// check if the nonce is correct
	if req.Nonce != senderAccount.Nonce {
		return errors.New("")
	}

	// check Credit
	if numRequest == 0 {
		copy(o.InitState, o.State)
	}

	if req.Nonce == 0 {
		senderAccount.Data.SetBytes(req.Data)
	}

	reqDataElement := fr.Element{}
	reqDataElement.SetBytes(req.Data)
	if reqDataElement.Equal(&senderAccount.Data) {
		senderAccount.Credit++
	} else {

		creditSub := uint64(math.Ceil(math.Log10(float64(senderAccount.Nonce)+1) / math.Log10(5)))

		if senderAccount.Credit < creditSub {
			senderAccount.Credit = 0
		} else {
			senderAccount.Credit = senderAccount.Credit - creditSub
		}
	}

	// update the nonce of the sender
	senderAccount.Nonce++

	// update the state of the operator
	copy(o.State[int(posSender)*SizeAccount:], senderAccount.Serialize())

	o.witnesses.SenderAccountsAfter[numRequest].Index = senderAccount.Index
	o.witnesses.SenderAccountsAfter[numRequest].Nonce = senderAccount.Nonce

	o.h.Reset()
	_, _ = o.h.Write(senderAccount.Serialize())
	bufSender := o.h.Sum([]byte{})
	copy(o.HashState[int(posSender)*o.h.Size():(int(posSender)+1)*o.h.Size()], bufSender)

	//  Set witnesses for the proof of inclusion of sender and receivers account after update
	buf.Reset()
	_, err = buf.Write(o.HashState)
	if err != nil {
		return err
	}
	merkleRootAfter, proofInclusionSenderAfter, _, err := merkletree.BuildReaderProof(&buf, o.h, segmentSize, posSender)
	if err != nil {
		return err
	}
	merkleProofHelperSenderAfter := merkle.GenerateProofHelper(proofInclusionSenderAfter, posSender, numLeaves)

	o.witnesses.RootHashesAfter[numRequest] = merkleRootAfter
	for i := 0; i < len(proofInclusionSenderAfter); i++ {
		o.witnesses.MerkleProofsSenderAfter[numRequest][i] = proofInclusionSenderAfter[i]

		if i < len(proofInclusionSenderAfter)-1 {
			o.witnesses.MerkleProofHelperSenderAfter[numRequest][i] = merkleProofHelperSenderAfter[i]
		}
	}

	return nil
}

func getAccountsFromState(state []byte) {

}

func (o *Operator) AttackDetect() {

	// i be index of account
	for i := 0; i < o.nbAccounts; i++ {

		rawCredit := binary.BigEndian.Uint64(o.InitState[int(i)*SizeAccount : int(i)*SizeAccount+8])
		rawNonce := binary.BigEndian.Uint64(o.InitState[int(i)*SizeAccount+16 : int(i)*SizeAccount+24])

		credit := binary.BigEndian.Uint64(o.State[int(i)*SizeAccount : int(i)*SizeAccount+8])
		nonce := binary.BigEndian.Uint64(o.State[int(i)*SizeAccount+16 : int(i)*SizeAccount+24])

		if nonce == 0 || nonce-rawNonce == 0 {
			continue
		}

		var creditRate = float64(credit) / float64(nonce)
		var gradient = (float64(credit) - float64(rawCredit)) / (float64(nonce) - float64(rawNonce))

		log.Println("Index:", i, "\t Gradient:", strconv.FormatFloat(gradient, 'f', 3, 64), "\t Nonce:", nonce, "\t Credit:", credit, "\t Rate:", strconv.FormatFloat(creditRate, 'f', 3, 64))

		filePath := "/mnt/e/godev/inftychain/rollup/groth16_attack.csv"
		strWrite := strconv.Itoa(i) + "," +
			strconv.Itoa(int(nonce)) + "," +
			strconv.FormatFloat(creditRate, 'f', 3, 64) + "," +
			strconv.FormatFloat(gradient, 'f', 3, 64) + "\n"

		file, err := os.OpenFile(filePath, os.O_WRONLY|os.O_APPEND, 0666)
		if err != nil {
			log.Println("open csv failed.", err)
		}
		defer file.Close()
		file.WriteString(strWrite)
	}
}
